from Core import Algorithm
import Core.Gadget as Gadget
import Core.IO as IO
import pandas as pd
import numpy as np
import datetime
import copy

def AverageCorrelation(database):
    #
    instruments2 = database.findWithFilter("Instruments", "Stock", {})
    instrumentList = database.findWithFilter("Instruments", "InstrumentList", {"Symbol": "000300.SH"},
                                             [("StdDateTime", pymongo.ASCENDING)])
    # for i in range(instruments2.__len__()):
    #    print(str(i) + " " + instruments2[i]["Symbol"])

    # instruments = []
    # for i in range(3):
    #    instruments.append(instruments2[i])
    # instruments = instruments2

    '''
    for instrument in instruments2:
        if instrument["Symbol"] == "000002.SZ":
            instruments.append(instrument)
        #if instrument["Symbol"] == "601069.SH":
        #    instruments.append(instrument)
        if instrument["Symbol"] == "002396.SZ":
            instruments.append(instrument)
        #if instrument["Symbol"] == "601098.SH":
        #    instruments.append(instrument)
    '''

    datetime1 = datetime(2005, 1, 1)
    datetime1 = ToUTCDateTime(datetime1)
    datetime1 = ToLocalDateTime(datetime1)
    datetime2 = datetime(2016, 9, 20)
    bmBarSeries = database.getDataSeries("000300.SH_Time_86400_Bar", datetime1, datetime2)

    '''
    barsCollectionBySymbol = {}
    for instrument in instruments:
        symbol = instrument["Symbol"]
        print("Load DataSeries " + symbol)
        barSeries = database.getDataSeries(symbol + "_Time_86400_Bar",datetime1,datetime2)
        barsCollectionBySymbol[symbol] = barSeries
    '''

    datetime2 = datetime1
    for bar in bmBarSeries:  # Loop Days

        if bar["StdDateTime"].month != datetime2.month:  # MonthChange
            datetime2 = bar["StdDateTime"]
            print(ToDateTimeString(datetime1) + " To " + ToDateTimeString(datetime2))

            # ---Make Sure HS300 List By DateTime---
            instruments = None
            if datetime1 < instrumentList[0]["StdDateTime"]:
                instruments = instrumentList[0]
            elif datetime1 > instrumentList[instrumentList.__len__() - 1]["StdDateTime"]:
                instruments = instrumentList[instrumentList.__len__() - 1]
            else:
                for i in range(1, instrumentList.__len__()):
                    if datetime1 > instrumentList[i - 1]["StdDateTime"] and datetime1 < instrumentList[i][
                        "StdDateTime"]:
                        instruments = instrumentList[i - 1]
                        break

            instruments = instruments["Values"]
            ###
            correlations = []
            countCorr = 0
            for i in range(instruments.__len__()):
                symbol1 = instruments[i]["Symbol"]
                # print(ToDateTimeString(datetime1) + " i " + str(i) + " " + symbol1)
                # barSeries1 = barsCollectionBySymbol[symbol1]
                barSeries1 = database.getDataSeries(symbol1 + "_Time_86400_Bar", datetime1, datetime2)

                if barSeries1.count() == 0:
                    continue
                if barSeries1[0]["StdDateTime"] > datetime2:
                    continue

                index11 = barSeries1.getIndex(datetime1)
                index12 = barSeries1.getIndex(datetime2)
                count = index12 - index11 + 1
                tradeCount1 = 0
                for day in range(index11, index12 + 1):
                    if barSeries1[day]["TradeStatus"] != "停牌一天":
                        tradeCount1 = tradeCount1 + 1
                if tradeCount1 / count < 0.5:
                    continue

                for j in range(instruments.__len__()):
                    countCorr = countCorr + 1
                    if j <= i:
                        continue

                    symbol2 = instruments[j]["Symbol"]

                    print("Correlate i " + str(i) + " " + symbol1 + " To " + symbol2 + " " + ToDateTimeString(
                        datetime1) + " j " + str(j) + " " + symbol2)
                    # barSeries2 = barsCollectionBySymbol[symbol2]
                    barSeries2 = database.getDataSeries(symbol2 + "_Time_86400_Bar", datetime1, datetime2)

                    if barSeries2.count() == 0:
                        print("j No Data")
                        continue
                    if barSeries2[0]["StdDateTime"] > datetime2:
                        continue

                    corr = Algorithm.CalcCorrelation(database, barSeries1, barSeries2)
                    # corr = 0
                    if corr != None:
                        correlations.append(corr)
                        data = {}
                        data["StdDateTime"] = datetime2
                        data["Symbol1"] = symbol1
                        data["Symbol2"] = symbol2
                        data["Corr"] = corr
                        data["Key"] = symbol1 + "_" + symbol2 + "_" + ToDateTimeString(datetime2)
                        database.upsert("Test", "Correlation2", {"Key": data["Key"]}, data)

                    print(ToDateTimeString(datetime2) + " Correlate [" + symbol1 + "] To [" + symbol2 + "] : " + str(
                        corr) + " countCorr " + str(countCorr))
                    kkwood = 1

            ave = np.mean(correlations)
            print(ToDateTimeString(datetime2) + " # " + str(correlations.__len__()) + " Ave Correlation " + str(ave))
            output = []
            output.append(ToDateTimeString(datetime2))
            output.append(correlations.__len__())
            output.append(ave)
            AppendListToFile("G:/TempData/AveCorrelation4.csv", output)
            datetime1 = bar["StdDateTime"]


def AveCorrelation2(database):

    # ---Build Baseline---
    datetime1 = datetime.datetime(2015, 12, 1)
    datetime2 = datetime.datetime(2019, 4, 10)
    datetime1 = Gadget.ToUTCDateTime(datetime1)
    datetime2 = Gadget.ToUTCDateTime(datetime2)

    instruments = [{"Symbol":"000001.SH"}]
    df = IO.LoadMultiInstrumentsFactorsAsDataFrame(database, "DailyReturn", datetime1, datetime2, instruments)

    # ---load Valid Symbols---
    validSymbols = []
    instruments = database.Find("Instruments", "Stock")
    for instrument in instruments:
        validSymbols.append(instrument["Symbol"])

    # ---Load Total---
    filter = {}
    gte = {"Value": {"$gte": -0.11}}
    lte = {"Value": {"$lte": 0.11}}
    filter["$and"] = [gte, lte]
    filter["$and"].append({"StdDateTime": {"$gte": datetime1}})
    filter["$and"].append({"StdDateTime": {"$lte": datetime2}})
    #filter["$or"] = []
    #filter["$or"].append({"Symbol": "000683.SZ"})
    #filter["$or"].append({"Symbol": "000685.SZ"})
    #filter["$or"].append({"Symbol": "000686.SZ"})
    #filter["$or"].append({"Symbol": "000697.SZ"})
    #filter["$or"].append({"Symbol": "000698.SZ"})

    #
    print("Load Return Factor Values", datetime.datetime.now())
    returns = database.Find("Factor", "DailyReturn", filter)

    # ---Group by Symbols---
    print("Format Returns by Symbol", datetime.datetime.now())
    returnsBySymbol = {}
    for r in returns:
        symbol = r["Symbol"]
        stdDatetime = r["StdDateTime"]
        value = r["Value"]
        if symbol not in validSymbols:
            continue
        if symbol not in returnsBySymbol.keys():
            returnsBySymbol[symbol] = []
        #
        returnsBySymbol[symbol].append([stdDatetime, value])

    #
    def FullBatch(dfBatchs, dfBatch, batchCount):
        print("Full Batch", batchCount)
        dfBatch.drop(columns=["000001.SH"], inplace=True)
        dfBatchs.append(dfBatch)
        dfBatch.to_csv("d:/Data/AveCorrelation/Batch_" + str(batchCount) + ".csv")

    # ---Merge---
    print("Merge", datetime.datetime.now())
    batchSize = 100
    dfBatchs = []
    i = 0
    dfBatch = copy.deepcopy(df)
    for symbol, returns in returnsBySymbol.items():
        i += 1
        #if symbol == "002334.SZ":
        #    print("Merging", symbol, i)

        # Create Dataframe
        tempDf = pd.DataFrame(returns, columns=["StdDateTime", symbol])
        dfBatch = pd.merge(dfBatch, tempDf, on='StdDateTime', how='left')

        # Full, Move to Next
        if (dfBatch.shape[1] - 2) >= batchSize:
            batchCount = int(i / batchSize)
            FullBatch(dfBatchs, dfBatch, batchCount)
            dfBatch = copy.deepcopy(df)

    # ---Dont forget last batch---
    batchCount = int(i / batchSize) + 1
    FullBatch(dfBatchs, dfBatch, batchCount)

    # ---Merge Batch Togather---
    print("Merge Batch", datetime.datetime.now())
    for dfBatch in dfBatchs:
        names = dfBatch.columns.values
        df = pd.merge(df, dfBatch, on='StdDateTime', how='left')

    # ---Checking missing---
    names = df.columns.values
    count = 0
    for symbol, returns in returnsBySymbol.items():
        if symbol not in names:
            print(symbol)
            count += 1

    # ---
    print("total columns", df.shape[1], " rows ", df.shape[0])

    # ---Check Na--
    def CheckNa(df, threhold):
        toDelete = []
        names = df.columns.values
        for name in names:
            # print(name)
            countNa = df[name].isna().sum()
            if countNa > threhold:
                # print(name, "#Na", countNa)
                toDelete.append(name)
        return toDelete

    # ---Calc Correlation---
    print("Calc Correlations", datetime.datetime.now())
    newData = []
    windowDays = 20
    l = len(df)
    for dayi in range(windowDays, l):
        stdDatetime = df["StdDateTime"][dayi]
        sDate = stdDatetime.strftime('%Y-%m-%d')
        #print(dayi, stdDatetime)

        # ---Get Slice---
        tempDf = df[dayi-windowDays:dayi]
        tempDf.drop(columns=["000001.SH"], inplace=True)

        # ---Drop Missing---
        # toDelete = CheckNa(tempDf, 0)
        # tempDf = tempDf.drop(columns=toDelete)

        # print(temp)
        # corrmat = tempDf.corr().abs()
        corrmat = tempDf.corr()
        # corrmat.to_csv("d:/Data/AveCorrelation/Correlation_" + str(sDate) + ".csv")

        print("Calced Corr Matrix", corrmat.shape[1], " rows ", corrmat.shape[0])
        sum = corrmat.sum()
        countZero = 0
        totalSum = 0
        for s in sum:
            if s == 0:
                countZero += 1
            totalSum += s
        validSymbol = corrmat.shape[0] - countZero
        finalSum = (totalSum - corrmat.shape[0] + countZero) * 0.5
        validCorrelation = validSymbol * (validSymbol - 1) * 0.5
        aveCorrelation = finalSum / validCorrelation

        # ---Calc Average Correlation---
        #correlations = []
        #for i in range(corrmat.shape[0]):
        #    for j in range(corrmat.shape[0]):
        #        if j > i:
        #            value = corrmat.iat[i,j]
        #            if not np.isnan(value):
        #                correlations.append(value)
        #
        #aveCorrelation = np.mean(correlations)
        print(dayi, stdDatetime, aveCorrelation, datetime.datetime.now())
        newData.append([sDate, aveCorrelation, validCorrelation, validSymbol])

    #
    dfIndex = pd.DataFrame(newData, columns=["StdDateTime", "AveCorrelation", "Valid", "Count"])
    dfIndex.to_csv("D:/Data/AveCorrelation/" + "AveCorrelationIndex.csv")
    pass


# ---Connect to DataBase, Find Series 连接数据库---
from Core.Config import *
cfgPathFilename = os.getcwd() + "/../config.json"
config = Config(cfgPathFilename)
database = config.DataBase("Mongo")
realtime = config.RealTime()

#
AveCorrelation2(database)

pass